package modules

import (
	"context"
	"fastdp-ops/cmd/agent/zaplog"
	pb "fastdp-ops/proto/exec"
	"fmt"
	"go.uber.org/zap"
	"sync"
)

type AgentExecServer struct {
	pb.UnimplementedExecServiceServer
}
type Module interface {
	// Run 在目标主机上执行模块逻辑（通过 SSH 会话）
	Run(req *pb.Execrequest) (*pb.ExecResponse, error)
}

var (
	registry = make(map[string]func() Module)
	mu       sync.RWMutex // 保护 registry 的并发安全
)

func (s *AgentExecServer) Exec(ctx context.Context, req *pb.Execrequest) (*pb.ExecResponse, error) {
	zaplog.Logger.Info("收到模块执行请求",
		zap.String("machine_id", req.MachineId),
		zap.String("Module", req.Module),
		zap.Any("Parameters", req.Parameters),
		zap.String("TaskId", req.TaskId),
	)
	mod, err := GetModule(req.Module)
	if err != nil {
		errResp := &pb.ExecResponse{
			MachineId: req.MachineId,
			TaskId:    req.TaskId,
			Success:   false,
			Error: &pb.ErrorDetail{
				Code:    1001,
				Message: fmt.Sprintf("获取模块失败：%s", err.Error()),
			},
		}
		return errResp, err
	}
	execResp, err := mod.Run(req)
	return execResp, err
}
func (s *AgentExecServer) BatchExec(batchreq *pb.BatchExecrequest, stream pb.ExecService_BatchExecServer) error {

	if batchreq.Ordered {
		for _, req := range batchreq.Tasks {
			zaplog.Logger.Info("收到批量模块执行请求",
				zap.String("machine_id", batchreq.BatchId),
				zap.String("machine_id", req.MachineId),
				zap.String("Module", req.Module),
				zap.Any("Parameters", req.Parameters),
				zap.String("TaskId", req.TaskId),
			)
			mod, err := GetModule(req.Module)
			if err != nil {
				// 模块获取失败，生成错误响应
				errResp := &pb.ExecResponse{
					MachineId: req.MachineId,
					TaskId:    req.TaskId,
					Success:   false,
					Error: &pb.ErrorDetail{
						Code:    1001,
						Message: fmt.Sprintf("获取模块失败：%s", err.Error()),
					},
				}
				if err := stream.Send(errResp); err != nil {
					zaplog.Logger.Sugar().Errorf("顺序执行发送响应失败：%s", err.Error())
					return err // 发送失败直接终止批量任务
				}
			}
			execResp, err := mod.Run(req)
			if err != nil {
				// 模块执行失败（补充错误信息）
				execResp = &pb.ExecResponse{
					MachineId: req.MachineId,
					TaskId:    req.TaskId,
					Success:   false,
					Error: &pb.ErrorDetail{
						Code:    1002,
						Message: fmt.Sprintf("模块执行失败：%s", err.Error()),
					},
				}
			}
			// 发送响应
			if err := stream.Send(execResp); err != nil {
				zaplog.Logger.Sugar().Errorf("顺序执行发送响应失败：%s", err.Error())
				return err // 发送失败终止批量任务
			}

		}
	} else {
		var wg sync.WaitGroup
		var mu sync.Mutex // 保护 stream.Send 的线程安全
		var sendErr error // 记录发送响应时的错误

		for _, req := range batchreq.Tasks {
			wg.Add(1)
			go func(req *pb.Execrequest) {
				defer wg.Done()
				zaplog.Logger.Info("收到批量模块并行执行请求",
					zap.String("batch_id", batchreq.BatchId),
					zap.String("machine_id", req.MachineId),
					zap.String("module", req.Module),
					zap.String("task_id", req.TaskId),
				)

				mod, err := GetModule(req.Module)
				if err != nil {
					errResp := &pb.ExecResponse{
						MachineId: req.MachineId,
						TaskId:    req.TaskId,
						Success:   false,
						Error: &pb.ErrorDetail{
							Code:    1001,
							Message: fmt.Sprintf("获取模块失败：%s", err.Error()),
						},
					}
					mu.Lock()
					if sendErr == nil { // 只记录第一次发送错误
						sendErr = stream.Send(errResp)
					}
					mu.Unlock()
					return
				}
				// 执行模块
				execResp, err := mod.Run(req)
				if err != nil {
					execResp = &pb.ExecResponse{
						MachineId: req.MachineId,
						TaskId:    req.TaskId,
						Success:   false,
						Error: &pb.ErrorDetail{
							Code:    1002,
							Message: fmt.Sprintf("模块执行失败：%s", err.Error()),
						},
					}
				}
				// 线程安全地发送响应
				mu.Lock()
				if sendErr == nil { // 只记录第一次发送错误
					sendErr = stream.Send(execResp)
				}
				mu.Unlock()
			}(req)

		}
		// 等待所有并行任务完成
		wg.Wait()
		// 检查并行执行中的发送错误
		if sendErr != nil {
			zaplog.Logger.Sugar().Errorf("并行执行发送响应失败：%s", sendErr.Error())
			return sendErr
		}
	}

	return nil

}

// Register 注册模块（在 init 函数中调用，每个模块自己注册）
func Register(name string, factory func() Module) {
	mu.Lock()
	defer mu.Unlock()
	registry[name] = factory
}

// GetModule 通过模块名获取实例（主程序调用）
func GetModule(name string) (Module, error) {
	mu.RLock()
	defer mu.RUnlock()
	factory, ok := registry[name]
	if !ok {
		err := fmt.Errorf("未知模块: %s", name)
		return nil, err // 返回错误，不再执行 factory()
	}
	return factory(), nil // 调用构造函数创建实例
}
